# -*- coding: utf-8 -*-

from __future__ import print_function

from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence

import os
import sys
import random
import pickle
import logging
import numpy as np

random.seed(42)


def load(path, name):
    return pickle.load(open(os.path.join(path, name), 'rb'))


corpus_data_path = 'datasets/char_level'
pkl_data_path = 'custom_data/char_level'

# configure logging
logger = logging.getLogger(os.path.basename(sys.argv[0]))
logging.basicConfig(format='%(asctime)s: %(levelname)s: %(message)s')
logging.root.setLevel(level=logging.INFO)

# prepare corpus
sentences = LineSentence(os.path.join(corpus_data_path, 'corpus.txt'))
vocab = load(pkl_data_path, 'vocabulary.pkl')

# run model
model = Word2Vec(sentences, size=100, min_count=2, window=5, sg=1, iter=10)
weights = model.wv.syn0
d = dict([(k, v.index) for k, v in model.wv.vocab.items()])
emb = np.zeros(shape=(len(vocab)+2, 100), dtype='float32')

for w, i in vocab.items():
    if w not in d:
        continue
    print(d)
    emb[i, :] = weights[d[w], :]

np.save(open(os.path.join(pkl_data_path, 'atec_task1_100_dim.embeddings'), 'wb'), emb)
logger.info('saved to "atec_task1_100_dim.embeddings"')
